#include "pub_defs.h"
#include <iostream>
#include <mpi.h>
#include <unistd.h>

int main(int argc, char* argv[]) {
    if (argc!=3){
        std::cerr<<"usage ./optimized M N" <<std::endl;
        return 1;
    }
    int m = std::atoi(argv[1]);
    int n = std::atoi(argv[2]);

    MPI_Init(NULL, NULL);
    int rank, size;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);

    uint8_t* arr;
    uint8_t* ans_arr;
    double start_time;
    // localm是每个进程中最大的m值（在每个进程中这个值应该是一样的）
    const int localm = (m-1)/size + 1;
    // masterm是主进程中数组的m值。
    const int masterm = m-localm*(size-1);
    uint8_t* local_arr = new uint8_t[(localm+2)*(n+2)];
    if (rank == 0) {
        arr = new uint8_t[(m+2)*(n+2)];
        ans_arr = new uint8_t[(m+2)*(n+2)];

        read_binfile(arr, (m+2)*(n+2), "input.bin");
        read_binfile(ans_arr, (m+2)*(n+2), "output.bin");
        // 记录开始的时间。
        start_time = wall_time();

        // 分发数据
        for (int dest = 1; dest < size ; ++dest) {
            int start = (masterm+localm*(dest-1))*(n+2);
            int end = (1+masterm+localm* dest +1)*(n+2);
            MPI_Send(&arr[start], end-start, MPI_UNSIGNED_CHAR, dest, 0, MPI_COMM_WORLD);
        }
        // 0号进程也就是master进程拷贝数据就可以了，不用发送数据。
        std::copy(&arr[1*(n+2)], &arr[(1+localm+1)*(n+2)], &local_arr[1*(n+2)]);
    }
    else {
        // 非0号进程接收数据。
        MPI_Recv(local_arr, (localm+2)*(n+2), MPI_UNSIGNED_CHAR, 0, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
    }
    // debug info
    // sleep(rank*1);
    // std::cout<< "rank: " << rank << std::endl;
    // print_arr(local_arr, localm+2, N+2);

    // 计算的时候才可以用的每个数组的m的大小。
    int array_m = rank==0?masterm:localm;

    gameoflife_optimized(local_arr, array_m, n,
                            rank, size);
    if (rank == 0) {
        std::copy(&local_arr[1*(n+2)], &local_arr[(1+localm)*(n+2)], &arr[1*(n+2)]);
        // 收拢数据
        for (int src = 1; src < size; ++src) {
            int start = (1+masterm+localm*(src-1))*(n+2);
            int end = (1+masterm+localm*src)*(n+2);
            MPI_Recv(&arr[start], end-start, MPI_UNSIGNED_CHAR, src, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
    }
    else{
        MPI_Send(&local_arr[1*(n+2)], localm*(n+2), MPI_UNSIGNED_CHAR, 0, 0, MPI_COMM_WORLD);
    }
    delete[] local_arr;
    if (rank == 0){
        double last_time = wall_time() - start_time;
        std::cout <<"M:"<<m<<" n: "<<n <<" size: "<< size<< " gen data baseline Time elapsed: " << last_time << " seconds." << std::endl;

        bool compre_result = compare_array(arr, ans_arr, (m+2)*(n+2));
        if (!compre_result) {
            std::cerr << "Error: output is not correct!" << std::endl;
            MPI_Finalize();
            delete[] arr;
            delete[] ans_arr;
            return 1;
        }
        std::cout << "Output is correct!" << std::endl;
        delete[] arr;
        delete[] ans_arr;
    }

    MPI_Finalize();
    return 0;
}